import json
from openai import OpenAI
from tqdm import tqdm  # 导入 tqdm 进度条库

# 配置 OpenAI 客户端
client = OpenAI(
    base_url='https://api.gpts.vin/v1',
    api_key='sk-cngkq0zvLkU5ucXBxmh8taDrMEs6XtFqexSh1WYZsZHYzZAM'  # 请替换为你的 API Key
)


# 调用大模型获取 JSON 格式的答案
def get_model_answer(question, choices):
    """
    发送问题和选项给大模型，获取 JSON 格式的答案
    """
    prompt = (
        f" you will handle the following medical question. Please think carefully about each option "
        f"Question: {question}\n"
        f"Options:\n" + "\n".join([f"{i}. {choice}" for i, choice in enumerate(choices)]) + "\n"
        f"Please provide your thought process and final answer in the following format:\n"
        f"{{\"question\": \"Question content\", \"think\": \"Your thought process\", \"answer\": \"Option index from 0 to length-1, corresponding to each choice\"}}\n"
    )

    try:
        completion = client.chat.completions.create(
            model="gpt-4o",
            messages=[{"role": "system", "content": "You are a professional doctor. When answering, please strictly adhere to the JSON format."},
                      {"role": "user", "content": prompt}]
        )
        response = completion.choices[0].message.content.strip()

        # 提取 JSON 部分
        start_idx = response.find("{")
        end_idx = response.rfind("}") + 1
        json_str = response[start_idx:end_idx]  # 获取有效的 JSON 字符串

        # 尝试解析 JSON
        return json.loads(json_str)
    except Exception as e:
        print(f"JSON 解析错误: {e}")
        return {"question": "未知", "think": "未知", "answer": ""}  # 如果解析失败，返回默认的错误信息


# 处理 JSONL 文件并即时保存模型返回的 JSON
def process_jsonl(input_file, output_file):
    with open(input_file, 'r', encoding='utf-8') as f:
        input_lines = f.readlines()

    try:
        with open(output_file, 'r', encoding='utf-8') as f_out:
            processed_lines = f_out.readlines()
            processed_count = len(processed_lines)
    except FileNotFoundError:
        processed_count = 0

    print(f"已处理问题数: {processed_count}，剩余待处理问题数: {len(input_lines) - processed_count}")

    with tqdm(total=len(input_lines), initial=processed_count, desc="处理问题", unit="问题") as pbar:
        for i, line in enumerate(input_lines):
            if i < processed_count:
                continue  # 跳过已处理的部分

            data = json.loads(line)
            question = data['question']
            choices = data['choices']

            model_response = get_model_answer(question, choices)

            if model_response:
                with open(output_file, 'a', encoding='utf-8') as f_out:
                    f_out.write(json.dumps(model_response, ensure_ascii=False) + "\n")

            pbar.update(1)

    print(f"模型返回的 JSON 结果已保存至 {output_file}")



# 评估正确率函数
def evaluate_accuracy(input_file, output_file):
    """
    评估模型输出文件的正确率。
    """
    with open(input_file, 'r', encoding='utf-8') as f_in, open(output_file, 'r', encoding='utf-8') as f_out:
        input_data = [json.loads(line) for line in f_in]  # 读取输入文件中的所有数据
        output_data = [json.loads(line) for line in f_out]  # 读取输出文件中的所有数据

        correct = 0
        total = len(input_data)

        for input_item, output_item in zip(input_data, output_data):
            correct_answer = input_item['answer']
            model_answer = output_item['answer']

            if str(correct_answer) == str(model_answer):  # 如果模型答案与正确答案相同
                correct += 1

        accuracy = (correct / total) * 100 if total > 0 else 0
        print(f"正确率: {accuracy:.2f}%")


# 主函数
def main():
    input_file = "G:\desktop\论文\data\mmlu\\professional_medicine\\test.jsonl"  # 你的输入 JSONL 文件路径
    output_file = "G:\desktop\论文\data\mmlu\\professional_medicine\\output_cot_4o.jsonl"  # 结果保存路径
    process_jsonl(input_file, output_file)

    # 评估模型输出的准确率
    evaluate_accuracy(input_file, output_file)


# 运行主函数
if __name__ == "__main__":
    main()
